import chromadb
import pandas as pd
import os
import time

# ==================== 用户配置区域 ====================

# 1. 公司ChromaDB的连接信息
#    (1)如果公司提供的是一个URL和端口号，用的是服务器模式(Client/Server), 请使用下面的配置（取消注释并填写 HttpClient 那一行代码）
#    并填写正确的主机(host)和端口(port),并且注释掉 PersistentClient和DB_PATH 那2行。
#    CLIENT = chromadb.HttpClient(host="your_company_chroma_host", port="8000")
#
#    (2)如果公司也是将数据库文件存在服务器的某个目录里，即文件持久化模式(Persistent), 请使用这个配置（取消注释并填写 PersistentClient和DB_PATH 那2行）
#    并填写数据库文件所在的路径
DB_PATH = "/path/to/your/company/vector_database"
CLIENT = chromadb.PersistentClient(path=DB_PATH)

# 2. 要导入的集合名称 (可以和之前一样，也可以是新的)
COLLECTION_NAME = "technical_standards"

# 3. 包含向量数据的 Parquet 文件路径
#exported_vectors.parquet 文件需要和 import_to_chroma.py 位于同一个目录下
PARQUET_FILE = "exported_vectors.parquet"

# 4. 每次批量导入的数据条数
#    可以根据服务器性能和内存大小调整。500-1000是比较合理的值。
BATCH_SIZE = 500

# ==================== 脚本主逻辑 ====================

def main():
    """主函数，执行数据导入流程"""
    
    print("--- 开始批量导入任务 ---")

    # 步骤 1: 检查 Parquet 文件是否存在
    if not os.path.exists(PARQUET_FILE):
        print(f"❌ 错误：找不到 Parquet 文件 '{PARQUET_FILE}'。请确保文件在当前目录下。")
        return

    # 步骤 2: 连接并准备ChromaDB集合
    try:
        print(f"🔗 正在连接到 ChromaDB...")
        # 使用 get_or_create_collection, 如果集合不存在会自动创建
        collection = CLIENT.get_or_create_collection(
            name=COLLECTION_NAME,
            metadata={"hnsw:space": "cosine"}  # 确保索引类型一致
        )
        print(f"✅ 成功获取或创建集合: '{COLLECTION_NAME}'")
        initial_count = collection.count()
        print(f"   - 导入前集合内有 {initial_count} 条数据。")
    except Exception as e:
        print(f"❌ 错误：无法连接或获取ChromaDB集合: {e}")
        return

    # 步骤 3: 读取 Parquet 文件
    try:
        print(f"📖 正在读取 Parquet 文件: '{PARQUET_FILE}'...")
        df = pd.read_parquet(PARQUET_FILE)
        total_rows = len(df)
        if total_rows == 0:
            print("⚠️ 文件中没有数据，任务结束。")
            return
        print(f"   - 文件中包含 {total_rows} 条数据。")
    except Exception as e:
        print(f"❌ 错误：读取 Parquet 文件失败: {e}")
        return

    # 步骤 4: 分批次导入数据
    total_batches = (total_rows - 1) // BATCH_SIZE + 1
    print(f"🚀 准备分 {total_batches} 个批次进行导入...")
    
    start_time = time.time()

    for i in range(0, total_rows, BATCH_SIZE):
        batch_num = (i // BATCH_SIZE) + 1
        end_idx = min(i + BATCH_SIZE, total_rows)
        
        # 获取当前批次的数据
        batch_df = df.iloc[i:end_idx]
        
        # 从DataFrame中提取ChromaDB需要的数据格式
        ids = batch_df["id"].tolist()
        embeddings = batch_df["embedding"].tolist()
        documents = batch_df["document"].tolist()
        metadatas = batch_df["metadata"].tolist()

        try:
            # 执行批量添加
            # ChromaDB的 .add() 方法是幂等的，如果ID已存在，它会更新数据
            collection.add(
                ids=ids,
                embeddings=embeddings,
                documents=documents,
                metadatas=metadatas
            )
            print(f"   - 批次 {batch_num}/{total_batches} 完成，成功导入 {len(ids)} 条数据。")
        except Exception as e:
            print(f"❌ 错误：处理批次 {batch_num} 时失败: {e}")
            # 可选择跳过此批次继续，或直接中断
            continue
            
    end_time = time.time()
    duration = end_time - start_time

    # 步骤 5: 最终验证
    print("\n--- 导入任务完成 ---")
    print(f"⏱️ 总耗时: {duration:.2f} 秒")
    try:
        final_count = collection.count()
        print(f"✅ 最终验证：集合 '{COLLECTION_NAME}' 中现在共有 {final_count} 条数据。")
        print(f"   - 本次任务新增/更新了 {final_count - initial_count} 条数据。")
    except Exception as e:
        print(f"⚠️ 无法获取最终统计信息: {e}")

if __name__ == "__main__":
    main()
